% 2-D, MULTI-GROUP, FIXED SOURCE DIFFUSION IN CARTESIAN GEOMETRY
% j.roberts, 1/10

clear; 
%clc; %clf;


disp('-------------------------------------------')
disp('--       beginning diffusion solver      --')
disp('-------------------------------------------')



%---- INPUT FILE (JUST AN "M" FILE) --------------------------------------=
twodbenchmark_input % 2-D IAEA benchmark, very close
%simple2d_input      % simple 1 box problem for reference
%simple2dRESP_input    % resp fct for 1 box problem
%twoelement_input % simple 2 box problem

%---- PROBLEM INPUT TEMPLATE ---------------------------------------------=
%   it:  0 for fixed source, 1 for eigenvalue
% it = 1;
%   basic data, to be replaced by cross-section routine
% dat = [ 1.500   0.030   0.000  0.000  0.000  
%         0.420   0.080   0.125  0.020  0.000
%         2.000   0.040   0.000  0.000  0.000
%         0.300   0.010   0.000  0.040  0.000 ];  
%   number of groups and number of materials
% numg = 2;  numm = 2;
%  define x and y coarse meshes
% xcm  = [   0   5    10       ];
% xfm  = [     20   20         ];
% ycm  = [   0   5   ];
% yfm  = [     20    ]*1;
%   coarse mesh material placement
% mt   = [  1  
%           2  ];  
%   boundary conditions
%   L,R,B,T denote left, right, bottom, top
%   possible values:
%       0 -- vacuum (zero entrant current)
%       1 -- reflective condition
%       2 -- incident current
% BCL = 1;  BCR = 1;  BCB = 1;  BCT = 1;        
%   coarse mesh volumetric source
% src(:,:,1) = [ 0  
%                0 ];            
% src(:,:,2) = [ 0 
%                0 ];           
%   incident current distribution (optional)
%   note: this is uniform w/in coarse mesh; an option
%         will later be added for DLP's
% IBSL = [1 1 1 1 1
%         0 0 0 0 0]';
% IBSR = [1 1 1 1 1
%         0 0 0 0 0]';
% IBSB = [1 1 1 1 1
%         0 0 0 0 0]';
% IBST = [1 1 1 1 1
%         0 0 0 0 0]';

tic

%---- GENERATE THE COEFFICIENTS -------------------------------------------
[AB, AL, AC, AR, AT, S, dx, dy, dv, dc, siga, sct] = ...
    twoDcoefMG( dat, abdat, numg, numm, xcm, xfm, ycm, yfm, mt, src, it, ....
    [BCL BCR BCB BCT], IBSL, IBSR, IBSB, IBST );
nusig = S;
%---- SOLVE----------------------------------------------------------------
N = sum(xfm); M = sum(yfm);
% pad end of AB,AL
AB = [AB' zeros(numg,N+1)]';
AL = [AL' zeros(numg,1)]';
% pad top of AR,AT
AT = [zeros(numg,N+1) AT']';
AR = [zeros(numg,1) AR']';
%phi = zeros( (N+1)*(M+1), numg );

if it==0 %-----------------------------------------------------------------
%     phi = zeros( (N+1)*(M+1), numg );
%     scsrc = zeros( (N+1)*(M+1), 1); 
%     phieps = 1e-10;  mxit = 200;  it = 0; phierr = 1; 
%     keff = FF/(AA+(LL+LR+LB+LT));
%     while (phierr > phieps && it < mxit)
%         phi0 = phi;
%         A = spdiags([AB(:,1) AL(:,1) AC(:,1) AR(:,1) AT(:,1)], ...
%             [-N-1 -1 0 1 N+1],(N+1)*(M+1),(N+1)*(M+1));
%         phi(:,1) = A \ ( S(:,1) +  1/keff * 0.125 * dv.*phi(:,2) );
%         for g = 2:numg
%             scsrc   = zeros( (N+1)*(M+1), 1);
%             for gg = 1:g-1
%                 scsrc(:,1) = scsrc(:,1) + sct(:,g,gg).*phi(:,gg);
%             end
%             A = spdiags([AB(:,g) AL(:,g) AC(:,g) AR(:,g) AT(:,g)], ...
%                 [-N-1 -1 0 1 N+1],(N+1)*(M+1),(N+1)*(M+1));
%             phi(:,g) = A \ ( S(:,g) + scsrc(:,1) );
%         end
%         phierr = norm( phi0 - phi );
%         it = it + 1;
%     end
else % eigenvalue
    disp('begin keff iteration')
    nusig = S;  % just rename it right away to avoid confusion
    %s       = S;
    s = ones((N+1)*(M+1),1)/(ycm(end)+xcm(end));  % uniform guess
    keff    = 1.0;    % guess the initial keff
    errK    = 1;      errS    = 1; 
    epsK    = 1e-4;   epsS = 1e-3;
    iter    = 0; itmx = 300;
    phi     = zeros((N+1)*(M+1),numg); 
    A1 = spdiags([AB(:,1) AL(:,1) AC(:,1) AR(:,1) AT(:,1)], ...
                    [-N-1 -1 0 1 N+1],(N+1)*(M+1),(N+1)*(M+1));
    A2 = spdiags([AB(:,2) AL(:,2) AC(:,2) AR(:,2) AT(:,2)], ...
                    [-N-1 -1 0 1 N+1],(N+1)*(M+1),(N+1)*(M+1));

    while  ( ( errK>epsK ) || (errS>epsS) ) && iter < itmx
        % set to solve 1st group flux
        s = s/sum(s);
        phi(:,1) = A1 \ (s(:,1)/keff) ;
        if numg > 1
            for g = 2:numg
                scsrc   = zeros( (N+1)*(M+1), 1);
                % compute scattering source
                for gg = 1:g-1 
                    scsrc(:,1) = scsrc(:,1) + sct(:,g,gg).*phi(:,gg);
                end
                phi(:,g) = A2 \ ( scsrc(:,1) );
            end   
        end
        sold = s; 
        kold = keff;   
        s(:,1) = 0; % reset fission source
        for g = 1:numg
            s(:,1) = s(:,1) + nusig(:,g).*phi(:,g);
        end         
        keff = sum(s.*dv)*kold/sum(sold.*dv);
        indx = find(s);
        errS = max( abs((s(indx)-sold(indx))./s(indx)) );
        errK = abs( (keff-kold)/keff );
        iter = iter + 1;  % number of iterations
        if ( mod(iter,20)== 0)
            disp([' iter = ',num2str(iter), ' keff = ',num2str(keff)])
            disp([' errK = ',num2str(errK), ' errS = ',num2str(errS)])
        end
    end
    disp([ 'final result: keff = ',num2str(keff)])
    disp([ '              iter = ',num2str(iter)])
end

toc

%[CRL CRR CRB CRT LL LR LB LT] = cresp(phi,dx,dy,dc,mt,xfm,yfm,numg);

phindx = ones( (N+1)*(M+1), 1 );
CN = length(xfm);
CM = length(yfm);
% find indices for actual flux (not outer region)
for i = 1:N+1
    for j = 1:M+1
        for ci = 1:CN
            if i  <= sum(xfm(1:ci)), mtx = ci; break; end
        end
        for ci = 1:CM
            if j  <= sum(yfm(1:ci)), mty = ci; break; end
        end
        k = i+(j-1)*( N+1 );
        if mt(mtx,mty)==5
            phindx(k)=0;
        end
    end
end
%phi = phi/(sum(sum(nusig.*phi))/sum(phindx.*dv));

for i = 1:N+1
    for j = 1:M+1
        k = i+(j-1)*( N+1 );
        indx(k,1)=i;
        indx(k,2)=j;
    end
end
% generate x and y vals
xint = zeros(N+1,1); 
yint = zeros(M+1,1);
for i = 2:N+1
    xint(i)=xint(i-1)+dx(i-1);
end
for i = 2:M+1
    yint(i)=yint(i-1)+dy(i-1);
end
phimap = zeros(N+1,M+1);  pooA = phimap;
for g = 1:numg
    for k = 1:(N+1)*(M+1)
        i = indx(k,1);
        j = indx(k,2); 
        phimap(i,j,g)=phi(k,g);
        pooA(i,j,g)=siga(k,g);
        srcmap(i,j,g)=S(k,g);
    end
end

for g = 1:numg
    figure(g)
    [cs,h] = contourf( xint,yint,phimap(:,:,g)' );
    title(['group ',num2str(g),' flux'])
    xlabel('x')
    ylabel('y')
    clabel(cs,h,'labelspacing',200,'fontsize',15)
end
% 
% 
% for g = 1:numg
%     hold on
%     figure(numg+1)
%     fig = semilogy(xint,phimap(:,1,g),'k--');
%     col = [rand(1) rand(1) rand(1)];
%     set(fig,'Color',col,'LineWidth',2);
%     lab(g,:)=(['group ',num2str(g)]);
% end
% hold off
% legend(lab,0)
% 
% if it == 0
% figure(numg+2)
% %fig = plot(xint,(CRB(:,1)-CRT(:,1))./(CRB(:,1)),'k');
% %fig = plot(xint,(phimap(:,1,1)-phimap(:,end,1))./(phimap(:,1,1)),'k');
% fig = plot(xint,CRB(:,1),'k',xint,CRT(:,1),'r--');
% set(fig,'LineWidth',2);
% title('bottom vs. top')
% xlabel('x')
% figure(numg+3)
% %fig = plot(yint,(CRL(:,1)-CRR(:,1))./(CRL(:,1)),'k');
% %fig = plot(yint,(phimap(1,:,1)-phimap(end,:,1))./(phimap(1,:,1)),'k');
% fig = plot(yint,CRL(:,1),'k',yint,CRR(:,1),'g--');
% set(fig,'LineWidth',2);
% title('left vs. right')
% xlabel('y')
% end

% resp = 0;
% if resp==1
% % determine legendre coefficients -----------------------------------------
% N = sum(xfm);    % max degree
% k = 0:N;    % interval
% maxord = 2; % maxorder to which we compute P's
% P = zeros(N,N+1);
% P(1,1:N+1) = 1; 
% P(2,1:N+1) = 1 - (2*k)/N;
% for i = 1:maxord %otherwise N-1
%     P(i+2,1:N+1) = ( (2*i+1)*(N-2*k).*P(i+1,:) - i*(N+i+1)*P(i,:) ) ./ ...
%         ((i+1)*(N-i)) ;
% end
% 
% for i = 0:maxord % fading factorial: FF(N,i) = N*(N-1)*...*(N-i+1) = N*(N-1)
%             %                          / (N-i)*(N-i-1)*... = F(N)/F(N-i);
%     a = (2*i+1);
%     b = factorial(N+i+1)/factorial(N);
%     c = factorial(N)/factorial(N-i);
%     w(i+1)=a*c/b;
% end
% x = xint;

%for ord = 1:maxord+1
%for G = 1:numg
%f(1,:) = CRB(:,G)'; f(2,:) = CRT(:,G)'; f(3,:) = CRR(:,G)';  f(4,:) = CRL(:,G)';
% now we want to expand f into fi
%ff(ord,G) = w(ord) * f(2,:) * P(ord,:)'
%ff(:,2) = w(2) * f(:,:) * P(2,:)';
%ff(:,3) = w(3) * f(:,:) * P(3,:)';  0.17957072167623   0.13630066539502
%ff(:,4) = w(4) * f(:,:) * P(4,:)';  0.08659697495400   0.21339407264467
%f1 = ff(side,1)*P(1,:);
%f2 = ff(side,1)*P(1,:) + ff(side,2)*P(2,:);
%f3 = ff(side,1)*P(1,:) + ff(side,2)*P(2,:) + ff(side,3)*P(3,:);
%f4 = ff(side,1)*P(1,:) + ff(side,2)*P(2,:) + ff(side,3)*P(3,:) + ff(side,4)*P(4,:);
%end
%end

%end

% recomputing keff via losses and gains
%FF = sum( phi(:,2).*nusig(:,2) );
%AA = sum( phi(:,1).*siga(:,1) ) + sum( phi(:,2).*siga(:,2) );
%disp([' FF = ',num2str(FF)])
%disp([' AA = ',num2str(AA)])
%disp([' LL = ',num2str(LL+LR+LB+LT)])
%kk = FF/(AA+(LL+LR+LB+LT)), keff
%ff = ff/norm(ff)
%IBSB = ff(1,:);  
%IBST = IBSB;
%IBSR = ff(3,:);


